import glob
import math
import os.path
import pickle

import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile as wav
from cut_wave import *

def getLipLimitIndex(start, end, fs, nfft):
    span = fs / nfft
    basicPos = int(start / fs * nfft)
    interval = end - start
    freqOffset = int(interval // span)
    down = basicPos
    up = basicPos + freqOffset
    return down, up

def cancelAxis():
    plt.axis('off')
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    plt.margins(0, 0)

def subAdjacent(Zxx):
    time = Zxx.shape[1]
    shape = Zxx.shape
    result = np.copy(Zxx[:,:shape[1]-1])
    for i in range(1, time):
        result[:, i - 1] = (Zxx[:, i] - Zxx[:, i - 1])
    result = np.abs(result)
    return result

def savePltPic(dirPath, wavPath):
    if not os.path.exists(dirPath):
        os.makedirs(dirPath)
    wavPath = wavPath.replace(os.sep, '/')
    wavName = wavPath.split('/')[-1]
    saveName = os.path.splitext(wavName)[0] + '.jpg'
    savePath = os.path.join(dirPath, saveName)
    plt.savefig(savePath)

def getInnerDir(wavPath):
    wavPath = wavPath.replace(os.sep, '/')
    temp = wavPath.split('/')
    innerDir = os.path.join(temp[-3], temp[-2])
    return innerDir

def GaussWeight(x, sigma):
    w = 1 / ((2 * np.pi) ** 0.5) * np.exp(-x*x/2/sigma/sigma)
    return w

def GaussSmoothZeroPad(var, lenNum, sigma):
    varNum = len(var)
    varBefore = np.zeros(lenNum)
    varAfter = np.zeros(lenNum)
    var = np.concatenate((varBefore, var, varAfter))
    mid = lenNum//2
    coreIndex = [i - mid for i in range(lenNum)]
    coreIndex = np.array(coreIndex)
    w = GaussWeight(coreIndex, sigma)
    # plt.plot(w)
    # plt.show()
    result = []

    for i in range(lenNum, varNum+lenNum):
        temp = np.dot(w,var[i-mid:i+lenNum-mid])
        result.append(temp)
    return np.array(result)

def AvgZeroPad(var, lenNum):
    varNum = len(var)
    varBefore = np.zeros(lenNum)
    varAfter = np.zeros(lenNum)
    var = np.concatenate((varBefore, var, varAfter))
    mid = lenNum//2
    coreIndex = [i - mid for i in range(lenNum)]
    coreIndex = np.array(coreIndex)
    w = [1/lenNum] * lenNum
    w = np.array(w)
    # plt.plot(w)
    # plt.show()
    result = []

    for i in range(lenNum, varNum+lenNum):
        temp = np.dot(w,var[i-mid:i+lenNum-mid])
        result.append(temp)
    return np.array(result)

def getActivation(sumArr, thresh):
    arrLen = sumArr.shape[0]
    res = np.zeros(arrLen)
    for i in range(arrLen):
        if sumArr[i] >= thresh:
            res[i] = 1
    return res

def getClipIndex(activation):
    arrLen = len(activation)
    resIndex = []
    start = -1
    end = -1
    for i in range(arrLen):
        val = activation[i]
        if val == 1:
            if start == -1:
                start = i
        else:
            if start != -1:
                end = i
                resIndex.append([start, end])
                start = -1
                end = -1
    if start!=-1:
        resIndex.append([start, arrLen-1])
    return resIndex

def getConcatIndex(clipIndex, winLen):
    clipLen = len(clipIndex)
    res = []
    start = -1
    end = -1
    for i in range(clipLen):
        cur = clipIndex[i]
        if start == -1:
            start = cur[0]
            end = cur[1]
        else:
            if(cur[0] - end <= winLen):
                end = cur[1]
            else:
                res.append([start, end])
                start = -1
                end = -1
    if start!=-1:
        res.append([start, end])
    return res

def getMaxLenIndex(multiIndex):
    res = [0, 0]
    for cur in multiIndex:
        if res[1] - res[0] < cur[1] - cur[0]:
            res = cur
    return res

def getWordFromPath(path):
    path = path.replace(os.sep, '/')
    word = path.split('/')[-1].split('_')[0]
    return word

def getFinalIndexRange(index, cropLen):
    half = cropLen // 2
    indexLen = index[1] - index[0]
    if indexLen > cropLen:
        mid = index[0] + half
    else:
        mid = (index[1] - index[0]) // 2
    if mid - half < 0:
        start = 0
        cropRight = cropLen - mid
    else:
        start = mid - half
        cropRight = cropLen - half
    end = mid + cropRight
    return [start, end]

def getMeanAndStdOfWordLen(lengthDict):
    res = {}
    for key in lengthDict.keys():
        vals = lengthDict[key]
        vals = np.array(vals)
        mean = np.mean(vals)
        std = np.std(vals)
        res[key] = [mean, std]
    return res


def getCropLenOfOneWord(meanStd, word):

    [mean, std] = meanStd[word]
    cropLen = mean + 2 * std
    cropLen = math.ceil(cropLen)
    return cropLen